"""
HDF5绘图管理模块
负责所有matplotlib绘图相关的功能，专为HDF5数据优化
"""

import numpy as np
import matplotlib.pyplot as plt
import config


class PlotManager:
    """绘图管理类，负责所有绘图操作"""
    
    @staticmethod
    def get_plot_grid_size(number_of_subplots):
        """
        根据子图数量获取网格大小
        
        Args:
            number_of_subplots: 子图数量
            
        Returns:
            tuple: (行数, 列数)
        """
        if number_of_subplots <= 1:
            return 2, 1
        if number_of_subplots == 2:
            return 2, 2
        if number_of_subplots == 3:
            return 2, 3
        if number_of_subplots == 4:
            return 2, 4
        if number_of_subplots in [5, 6]:
            return 4, 3
        if number_of_subplots in [7, 8]:
            return 4, 4
        if number_of_subplots in [9, 10, 11, 12]:
            return 6, 4
        return 8, 4
    
    @staticmethod
    def scale_time_axis(t, smpl_period, time_unit):
        """
        缩放时间轴
        
        Args:
            t: 时间数组
            smpl_period: 采样周期
            time_unit: 时间单位
            
        Returns:
            tuple: (缩放后的时间数组, 时间单位标签)
        """
        time_unit_scale = config.TIME_UNIT_SCALES.copy()
        time_unit_scale['[samples]'] = 1 / smpl_period if smpl_period > 0 else 1
        
        t_unit_label = time_unit if time_unit != '[samples]' else '[samples]'
        scale = time_unit_scale.get(time_unit, 1e6)  # 默认为微秒
        
        return t * scale, t_unit_label
    
    @staticmethod
    def plot_wave_separated(ax_amp, ax_phase, t_scaled, amp_data, phase_data, 
                           plot_mode, legend_label, ch_name, font_size, line_width, 
                           t_unit_label, log_y, p_min=None, p_max=None):
        """
        绘制波形图（分离的幅度和相位）
        
        Args:
            ax_amp: 幅度轴
            ax_phase: 相位轴
            t_scaled: 缩放后的时间数组
            amp_data: 幅度数据
            phase_data: 相位数据
            plot_mode: 绘图模式 ('A/P' 或 'I/Q')
            legend_label: 图例标签
            ch_name: 通道名称
            font_size: 字体大小
            line_width: 线宽
            t_unit_label: 时间单位标签
            log_y: 是否使用对数Y轴
            p_min: 相位最小值
            p_max: 相位最大值
        """
        if plot_mode == 'A/P':
            # 绘制幅度和相位
            ax_amp.plot(t_scaled, amp_data, linewidth=line_width, label=legend_label)
            ax_phase.plot(t_scaled, phase_data, linewidth=line_width, label=legend_label)
            ax_amp.set_ylabel(f'{ch_name} A [a.u.]', fontsize=font_size)
            ax_phase.set_ylabel(f'{ch_name} P [deg]', fontsize=font_size)
        else:  # I/Q
            # 从幅度和相位计算I和Q分量
            complex_data = amp_data * np.exp(1j * np.deg2rad(phase_data))
            ax_amp.plot(t_scaled, np.real(complex_data), linewidth=line_width, label=legend_label)
            ax_phase.plot(t_scaled, np.imag(complex_data), linewidth=line_width, label=legend_label)
            ax_amp.set_ylabel('I [a.u.]', fontsize=font_size)
            ax_phase.set_ylabel('Q [a.u.]', fontsize=font_size)
        
        # 设置对数Y轴
        if log_y:
            ax_amp.set_yscale('log')
        
        # 设置相位范围
        if p_min is not None and p_max is not None:
            try:
                ax_phase.set_ylim(float(p_min), float(p_max))
            except ValueError:
                pass
        
        # 设置图例和标签
        ax_amp.legend(prop={'size': font_size - 2})
        ax_phase.legend(prop={'size': font_size - 2})
        ax_amp.tick_params(axis='x', labelbottom=False)  # 隐藏顶部图的x轴标签
        ax_phase.set_xlabel(f'Time {t_unit_label}', fontsize=font_size)
    
    @staticmethod
    def plot_fft_separated(ax_fft_amp, ax_fft_phase, t, amp_data, phase_data, smpl_period, 
                          legend_label, ch_name, font_size, line_width, 
                          f_start=None, f_end=None):
        """
        绘制FFT图（分离的幅度和相位）
        
        Args:
            ax_fft_amp: FFT幅度轴
            ax_fft_phase: FFT相位轴
            t: 时间数组
            amp_data: 幅度数据
            phase_data: 相位数据
            smpl_period: 采样周期
            legend_label: 图例标签
            ch_name: 通道名称
            font_size: 字体大小
            line_width: 线宽
            f_start: 频率起始值
            f_end: 频率结束值
        """
        N_fft = len(t)
        if N_fft == 0:
            return
        
        # 计算频率轴
        xf = np.fft.fftfreq(N_fft, d=smpl_period)
        
        # 计算dA/A
        mean_amp = np.mean(amp_data)
        y_amp = (amp_data - mean_amp) / mean_amp if mean_amp != 0 else amp_data - mean_amp
        yf_amp = np.fft.fft(y_amp)
        
        # 绘制幅度FFT
        ax_fft_amp.plot(
            np.fft.fftshift(xf), 
            20 * np.log10(np.abs(np.fft.fftshift(yf_amp)) + 1e-12), 
            linewidth=line_width, 
            label=legend_label
        )
        ax_fft_amp.set_ylabel(f'{ch_name} dA/A [dB]', fontsize=font_size)
        
        # 计算dP
        y_phase = phase_data - np.mean(phase_data)
        yf_phase = np.fft.fft(y_phase)
        
        # 绘制相位FFT
        ax_fft_phase.plot(
            np.fft.fftshift(xf), 
            np.abs(np.fft.fftshift(yf_phase)), 
            linewidth=line_width, 
            label=legend_label
        )
        ax_fft_phase.set_ylabel(f'{ch_name} dP [deg]', fontsize=font_size)
        
        # 设置图例和标签
        ax_fft_amp.legend(prop={'size': font_size - 2})
        ax_fft_phase.legend(prop={'size': font_size - 2})
        ax_fft_amp.tick_params(axis='x', labelbottom=False)
        ax_fft_phase.set_xlabel('Frequency [Hz]', fontsize=font_size)
        
        # 设置频率范围
        if f_start is not None and f_end is not None:
            try:
                f_start_val = float(f_start)
                f_end_val = float(f_end)
                ax_fft_amp.set_xlim(f_start_val, f_end_val)
                ax_fft_phase.set_xlim(f_start_val, f_end_val)
            except ValueError:
                pass
    
    @staticmethod
    def plot_figure_content(axes, data_pairs, parameters):
        """
        绘制图形内容
        
        Args:
            axes: matplotlib轴数组
            data_pairs: 数据对列表，每个元素包含 {
                'name': 通道名称,
                'amp_data': 幅度数据,
                'phase_data': 相位数据,
                't': 时间数组
            }
            parameters: 绘图参数
        """
        font_size = parameters['font_size']
        line_width = parameters['line_width']
        plot_mode = parameters['plot_mode']
        wave_or_fft = parameters['wave_or_fft']
        log_y = parameters['log_y']
        time_unit = parameters['time_unit']
        sampling_freq = parameters['sampling_freq']
        
        # 计算采样周期
        smpl_period = 1.0 / sampling_freq if sampling_freq > 0 else 1.0
        
        num_pairs = len(data_pairs)
        npR, npC = PlotManager.get_plot_grid_size(num_pairs)
        
        for idx, data_pair in enumerate(data_pairs):
            ch_name_full = data_pair['name']
            amp_data = data_pair['amp_data']
            phase_data = data_pair['phase_data']
            t = data_pair['t']

            # 提取简称作为Y轴标签
            ch_name_short = ch_name_full.split(':')[-1]
            
            # 缩放时间轴
            t_scaled, t_unit_label = PlotManager.scale_time_axis(t, smpl_period, time_unit)
            
            # 计算子图位置
            row, col = PlotManager._calculate_subplot_position(idx, num_pairs, npC)
            
            if row + 1 >= npR:
                continue  # 边界检查
            
            ax_amp = axes[row, col]
            ax_phase = axes[row + 1, col]
            
            legend_label = ch_name_short
            
            if wave_or_fft == 'Wave':
                PlotManager.plot_wave_separated(
                    ax_amp, ax_phase, t_scaled, amp_data, phase_data,
                    plot_mode, legend_label, ch_name_short, font_size, line_width, 
                    t_unit_label, log_y, parameters.get('p_min'), parameters.get('p_max')
                )
            else:  # FFT
                PlotManager.plot_fft_separated(
                    ax_amp, ax_phase, t, amp_data, phase_data, smpl_period, 
                    legend_label, ch_name_short, font_size, line_width,
                    parameters.get('f_start'), parameters.get('f_end')
                )
    
    @staticmethod
    def _calculate_subplot_position(idx, num_pairs, npC):
        """
        计算子图位置
        
        Args:
            idx: 索引
            num_pairs: 数据对数量
            npC: 列数
            
        Returns:
            tuple: (行, 列)
        """
        if num_pairs <= 4:
            return 0, idx
        elif num_pairs <= 8:
            return (idx // npC) * 2, idx % npC
        elif num_pairs <= 12:
            return (idx // npC) * 2, idx % npC
        else:
            return (idx // npC) * 2, idx % npC
    
    @staticmethod
    def setup_figure_layout(fig, title):
        """
        设置图形布局
        
        Args:
            fig: matplotlib图形对象
            title: 标题
        """
        fig.suptitle(title, fontsize=12)
        fig.tight_layout(rect=[0, 0, 1, 0.95])
    
    @staticmethod
    def create_figure_title(file_name, waveform_index, sampling_rate):
        """
        创建图形标题
        
        Args:
            file_name: 文件名
            waveform_index: 波形索引
            sampling_rate: 采样率
            
        Returns:
            str: 图形标题
        """
        return f"{file_name} | Waveform: {waveform_index} | Sampling Rate: {sampling_rate} MHz" 